%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  multi-group diffusion solver for simple geometries,  j.roberts        %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% GEOMETRY                                                               %
%   1 == 1-D slab geometry                                               %
%   2 == 2-D geometry (not implemented)                                  %
%   3 == 1-D cylindrical (not implemented)                               %
%   4 == 1-D spherical (not implemented)                                 %
% BOUNDARY CODITIONS                                                     %
%   0 == vacuum condition (boundary flux = 0)                            %
%   1 == reflective condition                                            %
%   2 == zero incoming current using partial current (better than 0)     %
% SOURCE OPTIONS                                                         %
%   0 == fixed-source (volume source)                                    %
%   1 == eigenvalue (fission source) (solves via power iteration)        %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%clear; clc;

bc = 2;         % boundary condition
it = 0;         % source option
analytic = 0;   % analytic example

% material and cross-section parameters
% cross-section format:
%    for all materials
%        for all groups
%           read difc(m,g) siga(m,g) nusi(m,g) sigsg1->gi sigsg2->gi ...
% we limit the data only to downscattering

%--A 2D PROBLEM I WAS TRYING TO MODEL
% numg = 2; % number of groups
% numm = 5; % number of materials
% file = 'xsec_bench_2g';
% [difc,siga,nusi,sigs] = rdxsec(file,numm,numg);
% 
% % water |   f1   |   f2   |   f3   |   f4   | water %
% % 1.158   3.321    3.321     3.321    3.321   1.158 %
% 
% % mesh information
% coarse  =   [0.000  1.158  4.479  7.800  11.121  14.442  15.600];
% fine    =   [     20     20     20     20      20      20    ];
% % material placement
% matl    =   [ 1 2 4 4 2 1 ];
% % source in each coarse mesh
% source    = [ 0 2 0 2 0 
%               0 0 1 0 0 ]';

% Simple 1-D problem
numg = 1; % number of groups
numm = 3; % number of materials
file = 'xsec_1g';
[difc,siga,nusi,sigs] = rdxsec(file,numm,numg);

% |mat1|mat1|

% mesh information
coarse  =   [ 0      1     2     3     4      5     ];
fine    =   [    100  100   100   100   100    ]/10;
% material placement
matl    =   [     3     3    3    3     3      ];
% source in each coarse mesh
source    = [     1     1    1    1     1      ]';       
          
% do some error checking on # cm's, fm's, mat's
numc   = length(coarse)-1;                    % number of course meshes
numf   = sum(fine);                           % total number of fine meshes

% initialize the finemesh data values
dd = zeros(numf,numg);
aa = zeros(numf,numg);
nf = zeros(numf,numg);
sc = zeros(numf,numg,numg);
ss = zeros(numf,numg);
dx = zeros(numf,1);

j = 0;
for i = 1:numc
    for g = 1:numg
	    dd( (j+1):(j+fine(i)), g ) = difc( matl(i), g  ); 
	    aa( (j+1):(j+fine(i)), g ) = siga( matl(i), g  );    
	    nf( (j+1):(j+fine(i)), g ) = nusi( matl(i), g  ); 
        for k = (j+1):(j+fine(i))
            sc( k, g, : ) = sigs( matl(i), g, :  );
        end
        dx( (j+1):(j+fine(i)) ) = (coarse(i+1)-coarse(i))/fine(i);
        if it == 0
            ss( (j+1):(j+fine(i)), g  ) = source(i, g ); % acts as fixed source
        end
    end
    j = sum(fine(1:i));
end
if sum( dx >= sqrt(dd(:,end)./aa(:,end)) ) > 0
    disp('meshes are too large!') % want dx < L = sqrt(DiffCoef/SigA)
end

if bc == 0
    n = numf-1; % (# MESHES)-1 flux points
elseif bc == 1
    n = numf+1; % (# MESHES)+1 flux points
elseif bc == 2
    n = numf+1; % (# MESHES)+1 flux points
else
    disp('*** bc unrecognized or is not yet supported ***')
end

% build the coefficient matrix (could use tridiag!)

A = zeros( n,n,numg );

for g = 1:numg
    if bc == 0  
        ii=0;
        A(1,2,g)   = -dd(2,g)/dx(2,1);    
        A(1,1,g)   = -A(1,2,g) + dd(1,g)/dx(1,1) + 0.25*(aa(2,g)*dx(2,1)+aa(1,g)*dx(1,1));
        A(n,n-1,g) = -dd(n,g)/dx(n,1); 
        A(n,n,g)   = -A(n,n-1,g) + dd(n+1,g)/dx(n+1,1) + 0.25*(aa(n+1,g)*dx(n+1,1)+aa(n,g)*dx(n,1));
    elseif bc == 1
        ii=1;
        A(1,2,g)   = -dd(1,g)/dx(1,1);
        A(1,1,g)   = -A(1,2,g) + 0.25*(aa(2,g)*dx(2,1)+aa(1,g)*dx(1,1)) ;
        A(n,n-1,g) = -dd(n-1,g)/dx(n-1,1);
        A(n,n,g)   = -A(n,n-1,g) + 0.25*(aa(n-1,g)*dx(n-1,1)+aa(n-1,g)*dx(n-1,1));
    else % better vacuum approximation
        ii=1;
        A(1,2,g)   = -dd(1,g)/dx(1,1);
        A(1,1,g)   = -A(1,2,g) + 0.5;
        A(n,n-1,g) = -dd(n-1,g)/dx(n-1,1);
        A(n,n,g)   = -A(n,n-1,g) + 0.5;
    end
    for i = (2-ii):(n-1-ii)  % flux "1" to "100"
        A(i+ii,i+ii-1,g) = -dd(i,g)/dx(i,1);
        A(i+ii,i+ii+1,g) = -dd(i+1,g)/dx(i+1,1);    
        A(i+ii,i+ii,g)   = -(A(i+ii,i+1+ii,g) + A(i+ii,i+ii-1,g)) + ...
            0.5*(aa(i+1,g)*dx(i+1,1)+aa(i,g)*dx(i,1));
    end
end

if it == 1
    nusig = zeros(n,numg);
else
    s = zeros(n,1);
end
dxx = zeros(n,1);
sct = zeros(n,numg,numg);
for g = 1:numg
    if bc > 0 % set values for constants at boundaries
        if it == 1
            if bc == 1
                nusig(1,g) = 0.5*dx(1)*nf(1,g);
                nusig(n,g) = 0.5*dx(n-1)*nf(n-1,g);  
            else
                nusig(1,g) = 0.0;
                nusig(n,g) = 0.0;
            end
        else    
            if bc == 1
                s(1,g)     = 0.5*dx(1)*ss(1,g);      % what's the right coeff on dx??
                s(n,g)     = 0.5*dx(n-1)*ss(n-1,g);
            else
                s(1,g) = 0.0;
                s(n,g) = 0.0;
            end
        end
        dxx(1)   = 0.5*dx(1);
        dxx(n)   = 0.5*dx(n-1);
        sct(1,g,:)   = 0.5*dx(1)*sc(1,g,:);
        sct(n,g,:)   = 0.5*dx(n-1)*sc(n-1,g,:);
    end
    for i = (1):(n-ii*2) 
        if it == 1
            %nusig(i+ii,g) = 0.5*(nf(i+1,g)*dx(i+1,1)+nf(i,g)*dx(i,1));
            nusig(i+ii,g) = 0.25*(dx(i)+dx(i+1))*(nf(i+1,g)+nf(i,g));
        else
            s(i+ii,g)     = 0.25*(dx(i)+dx(i+1))*(ss(i,g)+ss(i+1,g));
        end
        sct(i+ii,g,:)   = 0.25*(dx(i)+dx(i+1))*(sc(i+1,g,:)+ sc(i,g,:));
    end
end
for i = 1:(n-ii*2)
    dxx(i+ii,:) = 0.5*(dx(i+1,:)+dx(i,:));
end
if it == 1
    sindx = find( sum(nusig(:,:)') ); % if no nusig, so source  
    s = 0.4*ones(n,1).*( sum(nusig(:,:)' > 0 ) )';  
    s = 0.01 * s / max(s);
end

%------------- make the x-vector for plotting------------------------
x = zeros(length(dx)+1,1);
for i = 2:1:length(dx)+1
	x(i) = x(i-1,1) + dx(i-1,1);
end

phi     = zeros(n,g);
scsrc   = zeros(n,1);

%------------- convert to tri-diag ----------------------------------
for gg = 1:numg
    AU(:,gg) = diag(A(:,:,gg),1);     % upper diagonal
    AL(:,gg) = diag(A(:,:,gg),-1);    % lower diagonal
    AD(:,gg) = diag(A(:,:,gg),0);     % central diagonal
end

if it == 1  % eigenvalue search -------------------------------------
    [k s phi] = keffit1_2( AU, AL, AD, dxx, nusig, sct, s, 1e-8, 1e-9 );
else        % fixed source, probably want function ------------------
    for g = 1:numg
        % compute scattering source
        for k = 1:n
            for gg = 1:g
                scsrc(k,1) = scsrc(k,1) + sct(k,g,gg)*phi(k,gg);
            end
        end
        % solve for phi_g
        phi(:,g) = A(:,:,1)\( s(:,g)+scsrc(:,1) );
    end  
end


if analytic == 1
    % ------------- analytic solution --------------------------------
    % -Df''+Ea*f = S, D = 1, Ea = 0.5, S = 2 across all
    xx = 0:.1:10;
    f = -.6794610552e-2*sinh(.7071067810*xx)+ ...
        .6794610552e-2*sinh(.7071067810*xx-7.071067810)+4.;
    % ----------------------------------------------------------------
    plot(xx,f,'b.')
    hold on
    % maybe want other solutions for comparison? hardcoded transport?
end

clear figures
% % plot multigroup
% for g = 1:numg
%     if bc == 0
%         f = [0 phi(:,g)' 0];
%     else
%         f = phi(:,g);
%     end
%     plot(x,f,'k');
%     axis([min(x) max(x) 0 1.25*max(max(phi))]);
%     xlabel('{x(cm)}')
%     ylabel('\phi(x)')
%     hold on
% end
figure(1)
if numg==2
    if bc == 0
        plot(x,[0 phi(:,1)' 0],'k',x,[0 phi(:,2)' 0],'r--')
        axis([min(x) max(x) 0 1.25*max(max(phi))]);
        xlabel('{x(cm)}')
        ylabel('\phi(x)')
    else
        plot(x,phi(:,1),'k',x,phi(:,2),'r--')
        axis([min(x) max(x) 0 1.25*max(max(phi))]);
        xlabel('{x(cm)}')
        ylabel('\phi(x)')
    end
elseif numg==1
    if bc == 0
        plot(x,[0 phi(:,1)' 0],'k')
        axis([min(x) max(x) 0 1.25*max(max(phi))]);
        xlabel('{x(cm)}')
        ylabel('\phi(x)')
    else
        plot(x,phi(:,1),'k')
        axis([min(x) max(x) 0 1.25*max(max(phi))]);
        xlabel('{x(cm)}')
        ylabel('\phi(x)')
    end
end